package first.group.ssh.keyexchange;

import first.group.ssh.CipherSuite;
import first.group.ssh.ConnectionCallback;
import first.group.ssh.PacketConnectStage;
import first.group.ssh.packet.KeyExchangeInitPacket;
import first.group.ssh.packet.NewKeysPacket;
import first.group.ssh.packet.Packet;
import first.group.ssh.version.Version;
import java.io.IOException;
import java.util.Random;

public class KeyExchanger extends PacketConnectStage {

    private final KeyExchangeMethod.Exchanger exchanger;
    private final ConnectionCallback connectionCallback;
    private boolean finished = false;

    public KeyExchanger(
            Version clientVersion,
            Version serverVersion,
            KeyExchangeInitPacket clientInitPacket,
            KeyExchangeInitPacket serverInitPacket,
            Random random,
            CipherSuite suite,
            ConnectionCallback connectionCallback,
            KeyExchangeMethod.Callback exchangeCallback
    ) {
        this.connectionCallback = connectionCallback;
        this.exchanger = suite.keyExchangeMethod().newInstance(
                clientVersion,
                serverVersion,
                clientInitPacket,
                serverInitPacket,
                random,
                connectionCallback,
                exchangeCallback
        );
    }

    @Override
    public void init() {
        exchanger.init();
    }

    public void prepareFinish() {
        finished = true;
        connectionCallback.appendPacket(new NewKeysPacket());
    }

    @Override
    public void read(Packet packet) throws IOException {
        if (!finished) {
            exchanger.read(packet);
        } else {
            assert (packet != null);
            if (!(packet instanceof NewKeysPacket)) {
                throw new IllegalArgumentException("Want NewKeysPacket, but get " + packet.getClass().getName());
            }
            connectionCallback.finish();
        }
    }

    @Override
    public boolean shouldRead() {
        return exchanger.shouldRead();
    }

}
